#include "tftp.h"
#define LINE_BUF_SIZE 1024
//函数声明
void get(char *remote_file, char *local_file);
void put(char *filename);

struct tftp_packet {
    ushort cmd;
    union {
        ushort code;
        ushort block;
        char filename[2];
    };
    char data[DATA_SIZE];
};

struct tftp_request {
    int size;
    struct sockaddr_in client;
    struct tftp_packet packet;
};
int sock;
struct sockaddr_in server;
socklen_t addr_len;
int blocksize = DATA_SIZE;

int main(int argc, char **argv) {
    FILE *record=fopen("record.txt","wr+");
    char cmd_line[LINE_BUF_SIZE];
    char *buf;
    char *arg;
    int i;
    char *local_file;
    int done = 0;
    char *server_ip;
    unsigned short port = SERVER_PORT;
    addr_len = sizeof(struct sockaddr_in);
    if (argc < 2) {
        printf("%s ip [port default 69]\n", argv[0]);
        return 0;
    }
    printf("    download：get+filename\n");
    printf("    upload：put+filename\n");
    printf("    exit:quit\n");
    server_ip = argv[1];
    if (argc > 2) {
        port = (unsigned short)atoi(argv[2]);
    }
    printf("tftp at %s:%d", server_ip, port);
   
    if ((sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP)) < 0) {
        printf("socket can not be created.\n");
        return 0;
    }

   
    server.sin_family = AF_INET;
    server.sin_port = htons(port);
    inet_pton(AF_INET, server_ip, &(server.sin_addr.s_addr));



    while (1) {
        printf(">> ");
        memset(cmd_line, 0, LINE_BUF_SIZE);
        buf = fgets(cmd_line, LINE_BUF_SIZE, stdin);
        if (buf == NULL) {
            printf("\nBye.\n");
            return 0;
        }

        arg = strtok(buf, " \t\n");
        fprintf(record,"%s ",arg);
        if (arg == NULL) {
            continue;
        }

        if (strcmp(arg, "get") == 0) {
            arg = strtok(NULL, " \t\n");
            
            fprintf(record,"%s\n",arg);
            fflush(record);
            local_file = strtok(NULL, " \t\n");
            if (arg == NULL) {
                printf("Error: missing target\n");
            }
            else {
                if (local_file == NULL)local_file = arg;
                get(arg, local_file);
            }
        }
        else if (strcmp(arg, "put") == 0) {
            arg = strtok(NULL, " \t\n");
            
            fprintf(record,"%s\n",arg);
            fflush(record);
            if (arg == NULL) {
                printf("Error: missing target\n");
            }
            else {
                put(arg);
            }
        }
        else if (strcmp(arg, "quit") == 0) {
            fclose(record);
            break;
        }
        else {
            printf("Unknown command.\n");
        }
    }
    return 0;
}

//下载
void get(char *remote_file, char *local_file) {
    struct tftp_packet snd_packet, rcv_packet;
    int next_block = 1;
    int recv_n;
    int total_bytes = 0;
    struct tftp_packet ack;
    struct sockaddr_in sender;

    int r_size = 0;
    int time_wait_data;
    ushort block = 1;

  
    snd_packet.cmd = htons(CMD_RRQ);
    sprintf(snd_packet.filename, "%s%c%s%c%d%c", remote_file, 0, "octet", 0, blocksize, 0);
    sendto(sock, &snd_packet, sizeof(struct tftp_packet), 0, (struct sockaddr*)&server, addr_len);
    FILE *fp = fopen(local_file, "w");
    if (fp == NULL) {
        printf("Create file \"%s\" error.\n", local_file);
        return;
    }

    //开始计时
    struct timespec get_start, get_end;
    clock_gettime(CLOCK_REALTIME, &get_start);
    //接收数据
    snd_packet.cmd = htons(CMD_ACK);
    do {
        for (time_wait_data = 0; time_wait_data < PKT_RCV_TIMEOUT * PKT_MAX_RXMT; time_wait_data += 10000) {
       
            r_size = recvfrom(sock, &rcv_packet, sizeof(struct tftp_packet), MSG_DONTWAIT,
                (struct sockaddr *)&sender,
                &addr_len);
            if (r_size > 0 && r_size < 4) {
                printf("Bad packet: r_size=%d\n", r_size);
            }
            if (r_size >= 4 && rcv_packet.cmd == htons(CMD_DATA) && rcv_packet.block == htons(block)) {
                printf("DATA: block=%d, data_size=%d\n", ntohs(rcv_packet.block), r_size - 4);
                snd_packet.block = rcv_packet.block;
                sendto(sock, &snd_packet, sizeof(struct tftp_packet), 0, (struct sockaddr*)&sender, addr_len);
                fwrite(rcv_packet.data, 1, r_size - 4, fp);
                break;
            }
            usleep(10000);
        }
        if (time_wait_data >= PKT_RCV_TIMEOUT * PKT_MAX_RXMT) {
            printf("Wait for DATA #%d timeout.\n", block);
            goto get_error;
        }
        block++;
    } while (r_size == blocksize + 4);
    clock_gettime(CLOCK_REALTIME, &get_end);
    double gettime = (get_end.tv_sec - get_start.tv_sec) + (get_end.tv_nsec - get_start.tv_nsec) / 1e9;
    printf("Total %lf (s)\n", gettime);
    printf("Get end.\n");
    printf("Average speed:%.0lf bps\n",block*4096/gettime);
    //printf("\nReceived %d bytes.\n", total_bytes);

get_error:
    fclose(fp);
}


// Upload
void put(char *filename) {
    struct sockaddr_in sender;
    struct tftp_packet rcv_packet, snd_packet;
    int r_size = 0;
    int time_wait_ack;

  
    snd_packet.cmd = htons(CMD_WRQ);
    sprintf(snd_packet.filename, "%s%c%s%c%d%c", filename, 0, "octet", 0, blocksize, 0);
    sendto(sock, &snd_packet, sizeof(struct tftp_packet), 0, (struct sockaddr*)&server, addr_len);

    for (time_wait_ack = 0; time_wait_ack < PKT_RCV_TIMEOUT; time_wait_ack += 20000) {
  
        r_size = recvfrom(sock, &rcv_packet, sizeof(struct tftp_packet), MSG_DONTWAIT,
            (struct sockaddr *)&sender,
            &addr_len);
        if (r_size > 0 && r_size < 4) {
            printf("Bad packet: r_size=%d\n", r_size);
        }
        if (r_size >= 4 && rcv_packet.cmd == htons(CMD_ACK) && rcv_packet.block == htons(0)) {
            break;
        }
        usleep(20000);
    }
    if (time_wait_ack >= PKT_RCV_TIMEOUT) {
        printf("Could not receive from server.\n");
        return;
    }

    FILE *fp = fopen(filename, "r");
    if (fp == NULL) {
        printf("File not exists!\n");
        return;
    }

    int s_size = 0;
    int rxmt;
    ushort block = 1;
    snd_packet.cmd = htons(CMD_DATA);

    //开始计时
    struct timespec put_start, put_end;
    clock_gettime(CLOCK_REALTIME, &put_start);
    // 发送数据
    do {
        memset(snd_packet.data, 0, sizeof(snd_packet.data));
        snd_packet.block = htons(block);
        s_size = fread(snd_packet.data, 1, blocksize, fp);

        for (rxmt = 0; rxmt < PKT_MAX_RXMT; rxmt++) {
            sendto(sock, &snd_packet, s_size + 4, 0, (struct sockaddr*)&sender, addr_len);
            printf("Send %d\n", block);

            for (time_wait_ack = 0; time_wait_ack < PKT_RCV_TIMEOUT; time_wait_ack += 20000) {
        
                r_size = recvfrom(sock, &rcv_packet, sizeof(struct tftp_packet), MSG_DONTWAIT,
                    (struct sockaddr *)&sender,
                    &addr_len);
                if (r_size > 0 && r_size < 4) {
                    printf("Bad packet: r_size=%d\n", r_size);
                }
                if (r_size >= 4 && rcv_packet.cmd == htons(CMD_ACK) && rcv_packet.block == htons(block)) {
                    break;
                }
                usleep(20000);
            }
            if (time_wait_ack < PKT_RCV_TIMEOUT) {
   
                break;
            }
            else {
                continue;
            }
        }
        if (rxmt >= PKT_MAX_RXMT) {
            printf("Could not receive from server.\n");
            return;
        }

        block++;
    } while (s_size == blocksize);
    clock_gettime(CLOCK_REALTIME, &put_end);
    double puttime = (put_end.tv_sec - put_start.tv_sec) + (put_end.tv_nsec - put_start.tv_nsec) / 1e9;
    printf("Total %lf (s)\n", puttime);
    printf("Put end\nAverage speed:%.0lf bps\n",block*4096/puttime);

put_error:
    fclose(fp);

    return;
}






